package tree;

/**
 * @ClassName AVLTree
 * @Description  学习 AVL树
 * @Author ZJX
 * @Date 2024/8/22 22:34
 * @Version 1.0
 */
public class AVLTree {

    // 定义AVL树节点
    static class TreeNode {
        public int val; // 节点的值
        public TreeNode left;  // 左子节点
        public TreeNode right; // 右子节点
        public TreeNode parent; // 父节点指针
        public int balanceFactor; // 平衡因子

        public TreeNode(int val) {
            this.val = val;
        }
    }

    public TreeNode root;  // AVL树的根节点

    /*
     * 插入一个新值到AVL树
     * 插入后需要通过旋转来保持树的平衡
     */
    public void insert(int val) {
        // 1. 二叉搜索树的方式插入
        if (root == null) {
            root = new TreeNode(val);
            return;
        }

        // 从根节点开始遍历寻找插入位置
        TreeNode cur = root;
        TreeNode parent = null;

        while (cur != null) {
            parent = cur;
            if (cur.val > val) {
                // 插入左子树
                cur = cur.left;
            } else if (cur.val < val) {
                // 插入右子树
                cur = cur.right;
            } else {
                // 值相等，不插入重复值
                return;
            }
        }

        // 创建新节点并连接到父节点
        TreeNode node = new TreeNode(val);
        node.parent = parent;

        // 根据父节点值决定插入左子树或右子树
        if (parent.val > node.val) {
            parent.left = node;
        } else {
            parent.right = node;
        }

        // 2. 插入后调整平衡
        balanceTreeAfterInsertion(node);
    }

    /*
     * 在插入节点后，从新插入的节点开始，向上逐级检查并调整树的平衡。
     * 它确保整个树在插入节点后仍然是平衡的。
     */
    private void balanceTreeAfterInsertion(TreeNode node) {
        TreeNode cur = node;

        while (cur != null) {
            // 更新当前节点的平衡因子
            updateBalanceFactor(cur);

            // 检查当前节点是否失衡
            if (cur.balanceFactor > 1 || cur.balanceFactor < -1) {
                // 失衡时进行旋转调整
                cur = balance(cur);
                break; // 一旦调整完成，退出
            }

            // 向上回溯到父节点
            cur = cur.parent;
        }
    }

    /*
     * 更新节点的平衡因子
     * 平衡因子 = 左子树高度 - 右子树高度
     */
    private void updateBalanceFactor(TreeNode node) {
        if (node != null) {
            node.balanceFactor = height(node.left) - height(node.right);
        }
    }

    // 递归计算树的高度
    private int height(TreeNode node) {
        if (node == null) {
            return 0;
        }
        return 1 + Math.max(height(node.left), height(node.right));
    }

    /*
     * 根据节点的平衡因子选择合适的旋转方式
     */
    private TreeNode balance(TreeNode node) {
        if (node.balanceFactor == -2) {
            // 左子树过高，进行右旋或左右旋
            if (node.left.balanceFactor <= 0) {
                // 左子树平衡因子 <= 0，右旋
                return rightRotate(node);
            } else {
                // 左子树平衡因子 > 0，先左旋再右旋
                return leftRightRotate(node);
            }
        } else if (node.balanceFactor == 2) {
            // 右子树过高，进行左旋或右左旋
            if (node.right.balanceFactor >= 0) {
                // 右子树平衡因子 >= 0，左旋
                return leftRotate(node);
            } else {
                // 右子树平衡因子 < 0，先右旋再左旋
                return rightLeftRotate(node);
            }
        }
        return node;
    }

    /*
     * 左旋操作
     */
    private TreeNode leftRotate(TreeNode parent) {
        TreeNode subR = parent.right; // 右子节点
        TreeNode subRL = subR.left;   // 右子节点的左子节点

        // 旋转：右子节点上升为新父节点
        subR.left = parent;
        parent.right = subRL;

        // 更新父节点引用
        if (subRL != null) {
            subRL.parent = parent;
        }

        // 更新父节点指针
        subR.parent = parent.parent;
        if (parent.parent == null) {
            root = subR;  // 如果旋转发生在根节点，更新根节点
        } else if (parent == parent.parent.left) {
            parent.parent.left = subR;
        } else {
            parent.parent.right = subR;
        }
        parent.parent = subR;

        // 更新平衡因子
        updateBalanceFactor(parent);
        updateBalanceFactor(subR);

        return subR;
    }

    /*
     * 右旋操作
     */
    private TreeNode rightRotate(TreeNode parent) {
        TreeNode subL = parent.left;  // 左子节点
        TreeNode subLR = subL.right;  // 左子节点的右子节点

        // 旋转：左子节点上升为新父节点
        subL.right = parent;
        parent.left = subLR;

        // 更新父节点引用
        if (subLR != null) {
            subLR.parent = parent;
        }

        // 更新父节点指针
        subL.parent = parent.parent;
        if (parent.parent == null) {
            root = subL;  // 如果旋转发生在根节点，更新根节点
        } else if (parent == parent.parent.left) {
            parent.parent.left = subL;
        } else {
            parent.parent.right = subL;
        }
        parent.parent = subL;

        // 更新平衡因子
        updateBalanceFactor(parent);
        updateBalanceFactor(subL);

        return subL;
    }

    /*
     * 左右旋操作
     */
    private TreeNode leftRightRotate(TreeNode parent) {
        parent.left = leftRotate(parent.left); // 先左旋子树
        return rightRotate(parent);            // 再右旋父节点
    }

    /*
     * 右左旋操作
     */
    private TreeNode rightLeftRotate(TreeNode parent) {
        parent.right = rightRotate(parent.right); // 先右旋子树
        return leftRotate(parent);                // 再左旋父节点
    }

    /*
     * 判断是否为AVL树
     */
    public Boolean isAVLTree() {
        // 1. 判断是否为二叉搜索树(BST)
        // 2. 判断其是否平衡
        return isBST(root, null, null) && isBalanced(root);
    }

    /*
     * 判断是否是二叉搜索树
     */
    private Boolean isBST(TreeNode node, Integer min, Integer max) {
        if (node == null) {
            return true; // 空节点自然是BST
        }

        // 当前节点值是否在允许的范围内
        if ((min != null && node.val <= min) || (max != null && node.val >= max)) {
            return false;
        }

        // 递归检查左子树和右子树
        return isBST(node.left, min, node.val) && isBST(node.right, node.val, max);
    }

    /*
     * 判断树是否平衡
     */
    private boolean isBalanced(TreeNode root) {
        if (root == null) {
            return true;
        }

        // 计算左右子树高度，检查是否平衡
        int leftHeight = getHeight(root.left);
        int rightHeight = getHeight(root.right);

        return Math.abs(leftHeight - rightHeight) <= 1
                && isBalanced(root.left) && isBalanced(root.right);
    }

    // 获取二叉树高度
    private int getHeight(TreeNode root) {
        if (root == null) {
            return 0;
        }

        return 1 + Math.max(getHeight(root.left), getHeight(root.right));
    }
}
